from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks


class StructBERT():
    def __init__(self):
        model_path = "/nfs1/huggingface/hub/nlp_structbert_sentiment-classification_chinese-tiny"
        self.semantic_cls = pipeline(Tasks.text_classification, model_path)



    def main(self,learnerDialogue):
        test_text = learnerDialogue["learnerDialogue"]

        prediction=self.semantic_cls(input=test_text)

        # {'scores': [0.5520888566970825, 0.4479111135005951], 'labels': ['负面', '正面']}
        attitudeStatus1=prediction["labels"][0]
        attitudeScore=prediction["scores"][0]
        if attitudeStatus1=="负面":
            attitudeStatus="2"
        else:
            attitudeStatus="1"
        prediction_={"attitudeStatus":attitudeStatus,"attitudeScore":attitudeScore}

        return  prediction_

if __name__ == '__main__':
    StructBERT()
